using System.Collections.Immutable;
using System.Text;
using HotChocolate.Types.Analyzers.Helpers;
using HotChocolate.Types.Analyzers.Inspectors;
using HotChocolate.Types.Analyzers.Models;
using Microsoft.CodeAnalysis;
using static HotChocolate.Types.Analyzers.Helpers.GeneratorUtils;

namespace HotChocolate.Types.Analyzers.FileBuilders;

public sealed class DataLoaderFileBuilder : IDisposable
{
    private StringBuilder _sb;
    private CodeWriter _writer;
    private bool _disposed;

    public DataLoaderFileBuilder()
    {
        _sb = PooledObjects.GetStringBuilder();
        _writer = new CodeWriter(_sb);
    }

    public void WriteHeader()
    {
        _writer.WriteIndentedLine("// <auto-generated/>");
        _writer.WriteLine();
        _writer.WriteIndentedLine("#nullable enable");
        _writer.WriteIndentedLine("#pragma warning disable");
        _writer.WriteLine();
        _writer.WriteIndentedLine("using System;");
        _writer.WriteIndentedLine("using System.Runtime.CompilerServices;");
        _writer.WriteIndentedLine("using Microsoft.Extensions.DependencyInjection;");
        _writer.WriteIndentedLine("using GreenDonut;");
        _writer.WriteLine();
    }

    public void WriteBeginNamespace(string ns)
    {
        _writer.WriteIndentedLine("namespace {0}", ns);
        _writer.WriteIndentedLine("{");
        _writer.IncreaseIndent();
    }

    public void WriteEndNamespace()
    {
        _writer.DecreaseIndent();
        _writer.WriteIndentedLine("}");
        _writer.WriteLine();
    }

    public void WriteDataLoaderInterface(
        string name,
        bool isPublic,
        DataLoaderKind kind,
        ITypeSymbol key,
        ITypeSymbol value)
    {
        _writer.WriteIndentedLine(
            "{0} interface {1}",
            isPublic
                ? "public"
                : "internal",
            name);
        _writer.IncreaseIndent();

        _writer.WriteIndentedLine(
            kind is DataLoaderKind.Group
                ? ": global::GreenDonut.IDataLoader<{0}, {1}[]>"
                : ": global::GreenDonut.IDataLoader<{0}, {1}>",
            key.ToFullyQualifiedWithNullRefQualifier(),
            value.ToFullyQualifiedWithNullRefQualifier());

        _writer.DecreaseIndent();
        _writer.WriteIndentedLine("{");
        _writer.WriteIndentedLine("}");
        _writer.WriteLine();
    }

    public void WriteBeginDataLoaderClass(
        string name,
        string interfaceName,
        bool isPublic,
        DataLoaderKind kind,
        ITypeSymbol key,
        ITypeSymbol value,
        bool withInterface)
    {
        _writer.WriteIndentedLine(
            "{0} sealed partial class {1}",
            isPublic
                ? "public"
                : "internal",
            name);
        _writer.IncreaseIndent();
        _writer.WriteIndentedLine(
            kind is DataLoaderKind.Group
                ? ": global::GreenDonut.DataLoaderBase<{0}, {1}[]>"
                : ": global::GreenDonut.DataLoaderBase<{0}, {1}>",
            key.ToFullyQualifiedWithNullRefQualifier(),
            value.ToFullyQualifiedWithNullRefQualifier());
        if (withInterface)
        {
            _writer.WriteIndentedLine(", {0}", interfaceName);
        }
        _writer.DecreaseIndent();
        _writer.WriteIndentedLine("{");
        _writer.IncreaseIndent();
    }

    public void WriteEndDataLoaderClass()
    {
        _writer.DecreaseIndent();
        _writer.WriteIndentedLine("}");
    }

    public void WriteDataLoaderConstructor(
        string name,
        DataLoaderKind kind,
        ITypeSymbol keyType,
        ITypeSymbol valueType,
        ImmutableArray<CacheLookup> lookupMethods)
    {
        _writer.WriteIndentedLine("private readonly global::System.IServiceProvider _services;");
        _writer.WriteLine();

        if (kind is DataLoaderKind.Batch or DataLoaderKind.Group)
        {
            _writer.WriteIndentedLine("public {0}(", name);

            using (_writer.IncreaseIndent())
            {
                _writer.WriteIndentedLine("global::System.IServiceProvider services,");
                _writer.WriteIndentedLine("global::GreenDonut.IBatchScheduler batchScheduler,");
                _writer.WriteIndentedLine("global::GreenDonut.DataLoaderOptions options)");
                _writer.WriteIndentedLine(": base(batchScheduler, options)");
            }
        }
        else
        {
            _writer.WriteIndentedLine("public {0}(", name);

            using (_writer.IncreaseIndent())
            {
                _writer.WriteIndentedLine("global::System.IServiceProvider services,");
                _writer.WriteIndentedLine("global::GreenDonut.DataLoaderOptions options)");
                _writer.WriteIndentedLine(": base(AutoBatchScheduler.Default, options)");
            }
        }

        _writer.WriteIndentedLine("{");

        using (_writer.IncreaseIndent())
        {
            _writer.WriteIndentedLine("_services = services ??");

            using (_writer.IncreaseIndent())
            {
                _writer.WriteIndentedLine("throw new global::System.ArgumentNullException(nameof(services));");
            }

            if (lookupMethods.Length > 0)
            {
                _writer.WriteLine();

                foreach (var lookup in lookupMethods)
                {
                    _writer.WriteIndentedLine(
                        "global::{0}",
                        WellKnownTypes.PromiseCacheObserver);

                    using (_writer.IncreaseIndent())
                    {
                        if (lookup.IsTransform)
                        {
                            _writer.WriteIndentedLine(
                                ".Create<{0}, {1}, {2}>({3}.{4}, this)",
                                keyType.ToFullyQualifiedWithNullRefQualifier(),
                                valueType.ToFullyQualifiedWithNullRefQualifier(),
                                lookup.Method.Parameters[0].Type.ToFullyQualifiedWithNullRefQualifier(),
                                lookup.Method.ContainingType.ToFullyQualifiedWithNullRefQualifier(),
                                lookup.Method.Name);
                        }
                        else
                        {
                            _writer.WriteIndentedLine(
                                ".Create<{0}, {1}>({2}.{3}, this)",
                                keyType.ToFullyQualifiedWithNullRefQualifier(),
                                valueType.ToFullyQualifiedWithNullRefQualifier(),
                                lookup.Method.ContainingType.ToFullyQualifiedWithNullRefQualifier(),
                                lookup.Method.Name);
                        }

                        _writer.WriteIndentedLine(".Accept(this);");
                    }
                }
            }
        }

        _writer.WriteIndentedLine("}");
    }

    public void WriteDataLoaderLoadMethod(
        string containingType,
        IMethodSymbol method,
        bool isScoped,
        DataLoaderKind kind,
        ITypeSymbol key,
        ITypeSymbol value,
        ImmutableArray<DataLoaderParameterInfo> parameters)
    {
        _writer.WriteIndentedLine(
            "protected override async global::{0} FetchAsync(",
            WellKnownTypes.ValueTask);

        using (_writer.IncreaseIndent())
        {
            _writer.WriteIndentedLine(
                "global::{0}<{1}> keys,",
                WellKnownTypes.ReadOnlyList,
                key.ToFullyQualifiedWithNullRefQualifier());
            _writer.WriteIndentedLine(
                "global::{0}<{1}<{2}>> results,",
                WellKnownTypes.Memory,
                WellKnownTypes.Result,
                kind is DataLoaderKind.Group
                    ? $"{value.ToClassNonNullableFullyQualifiedWithNullRefQualifier()}[]?"
                    : value.ToNullableFullyQualifiedWithNullRefQualifier());
            _writer.WriteIndentedLine(
                "global::{0}<{1}{2}> context,",
                WellKnownTypes.DataLoaderFetchContext,
                value.ToFullyQualifiedWithNullRefQualifier(),
                kind is DataLoaderKind.Group ? "[]" : string.Empty);
            _writer.WriteIndentedLine(
                "global::{0} ct)",
                WellKnownTypes.CancellationToken);
        }

        _writer.WriteIndentedLine("{");

        using (_writer.IncreaseIndent())
        {
            if (isScoped && parameters.Any(p => p.Kind == DataLoaderParameterKind.Service))
            {
                _writer.WriteIndentedLine("await using var scope = _services.CreateAsyncScope();");
            }

            foreach (var parameter in parameters)
            {
                if (parameter.Kind is DataLoaderParameterKind.Service)
                {
                    _writer.WriteIndentedLine(
                        "var {0} = {1}.GetRequiredService<{2}>();",
                        parameter.VariableName,
                        isScoped ? "scope.ServiceProvider" : "_services",
                        parameter.Type.ToFullyQualifiedWithNullRefQualifier());
                }
                else if (parameter.Kind is DataLoaderParameterKind.SelectorBuilder)
                {
                    _writer.WriteIndentedLine(
                        "var {0} = context.GetState<{1}>(\"{2}\")",
                        parameter.VariableName,
                        parameter.Type.ToFullyQualifiedWithNullRefQualifier(),
                        parameter.StateKey);
                    _writer.IncreaseIndent();
                    _writer.WriteIndentedLine(
                        "?? global::GreenDonut.Data.DefaultSelectorBuilder.Empty;");
                    _writer.DecreaseIndent();
                }
                else if (parameter.Kind is DataLoaderParameterKind.PredicateBuilder)
                {
                    _writer.WriteIndentedLine(
                        "var {0} = context.GetState<{1}>(\"{2}\")",
                        parameter.VariableName,
                        parameter.Type.ToFullyQualifiedWithNullRefQualifier(),
                        parameter.StateKey);
                    _writer.IncreaseIndent();
                    _writer.WriteIndentedLine(
                        "?? global::GreenDonut.Data.DefaultPredicateBuilder.Empty;");
                    _writer.DecreaseIndent();
                }
                else if (parameter.Kind is DataLoaderParameterKind.SortDefinition)
                {
                    _writer.WriteIndentedLine(
                        "var {0} = context.GetState<{1}>(\"{2}\")",
                        parameter.VariableName,
                        parameter.Type.ToFullyQualifiedWithNullRefQualifier(),
                        parameter.StateKey);
                    _writer.IncreaseIndent();
                    _writer.WriteIndentedLine(
                        "?? {0}.Empty;",
                        parameter.Type.ToFullyQualifiedWithNullRefQualifier());
                    _writer.DecreaseIndent();
                }
                else if (parameter.Kind is DataLoaderParameterKind.QueryContext)
                {
                    _writer.WriteIndentedLine(
                        "var {0}_selector = context.GetState<global::{1}>(\"{2}\")?.TryCompile<{3}>();",
                        parameter.VariableName,
                        WellKnownTypes.SelectorBuilder,
                        DataLoaderInfo.Selector,
                        ((INamedTypeSymbol)parameter.Type).TypeArguments[0].ToFullyQualifiedWithNullRefQualifier());
                    _writer.WriteIndentedLine(
                        "var {0}_predicate = context.GetState<global::{1}>(\"{2}\")?.TryCompile<{3}>();",
                        parameter.VariableName,
                        WellKnownTypes.PredicateBuilder,
                        DataLoaderInfo.Predicate,
                        ((INamedTypeSymbol)parameter.Type).TypeArguments[0].ToFullyQualifiedWithNullRefQualifier());
                    _writer.WriteIndentedLine(
                        "var {0}_sortDefinition = context.GetState<global::{1}<{2}>>(\"{3}\");",
                        parameter.VariableName,
                        WellKnownTypes.SortDefinition,
                        ((INamedTypeSymbol)parameter.Type).TypeArguments[0].ToFullyQualifiedWithNullRefQualifier(),
                        DataLoaderInfo.Sorting);
                    _writer.WriteLine();

                    _writer.WriteIndentedLine(
                        "var {0} = global::{1}<{2}>.Empty;",
                        parameter.VariableName,
                        WellKnownTypes.QueryContext,
                        ((INamedTypeSymbol)parameter.Type).TypeArguments[0].ToFullyQualifiedWithNullRefQualifier());
                    _writer.WriteIndentedLine(
                        "if({0}_selector is not null || {0}_predicate is not null || {0}_sortDefinition is not null)",
                        parameter.VariableName);
                    _writer.WriteIndentedLine("{");
                    _writer.IncreaseIndent();
                    _writer.WriteIndentedLine(
                        "{0} = new global::{1}<{2}>({0}_selector, "
                        + "{0}_predicate, {0}_sortDefinition);",
                        parameter.VariableName,
                        WellKnownTypes.QueryContext,
                        ((INamedTypeSymbol)parameter.Type).TypeArguments[0].ToFullyQualifiedWithNullRefQualifier());
                    _writer.DecreaseIndent();
                    _writer.WriteIndentedLine("}");
                }
                else if (parameter.Kind is DataLoaderParameterKind.PagingArguments)
                {
                    _writer.WriteIndentedLine(
                        "var {0} = context.GetRequiredState<{1}>(\"{2}\");",
                        parameter.VariableName,
                        parameter.Type.ToFullyQualifiedWithNullRefQualifier(),
                        parameter.StateKey);
                }
                else if (parameter.Kind is DataLoaderParameterKind.ContextData)
                {
                    if (parameter.Parameter.HasExplicitDefaultValue)
                    {
                        var defaultValue = parameter.Parameter.ExplicitDefaultValue;
                        var defaultValueString = ConvertDefaultValueToString(defaultValue, parameter.Type);

                        _writer.WriteIndentedLine(
                            "var {0} = context.GetStateOrDefault<{1}>(\"{2}\", {3});",
                            parameter.VariableName,
                            parameter.Type.ToFullyQualifiedWithNullRefQualifier(),
                            parameter.StateKey,
                            defaultValueString);
                    }
                    else if (parameter.Type.IsNullableType())
                    {
                        _writer.WriteIndentedLine(
                            "var {0} = context.GetState<{1}>(\"{2}\");",
                            parameter.VariableName,
                            parameter.Type.ToFullyQualifiedWithNullRefQualifier(),
                            parameter.StateKey);
                    }
                    else
                    {
                        _writer.WriteIndentedLine(
                            "var {0} = context.GetRequiredState<{1}>(\"{2}\");",
                            parameter.VariableName,
                            parameter.Type.ToFullyQualifiedWithNullRefQualifier(),
                            parameter.StateKey);
                    }
                }
            }

            if (kind is DataLoaderKind.Cache)
            {
                _writer.WriteIndentedLine("for (var i = 0; i < keys.Count; i++)");
                _writer.WriteIndentedLine("{");

                using (_writer.IncreaseIndent())
                {
                    _writer.WriteIndentedLine("try");
                    _writer.WriteIndentedLine("{");

                    using (_writer.IncreaseIndent())
                    {
                        _writer.WriteIndentedLine("var key = keys[i];");
                        _writer.WriteIndented("var value = ");
                        WriteFetchCall(method, containingType, kind, parameters);
                        _writer.WriteIndentedLine(
                            "results.Span[i] = Result<{0}>.Resolve(value);",
                            value.ToNullableFullyQualifiedWithNullRefQualifier());
                    }

                    _writer.WriteIndentedLine("}");
                    _writer.WriteIndentedLine("catch (global::System.Exception ex)");
                    _writer.WriteIndentedLine("{");

                    using (_writer.IncreaseIndent())
                    {
                        _writer.WriteIndentedLine(
                            "results.Span[i] = Result<{0}>.Reject(ex);",
                            value.ToNullableFullyQualifiedWithNullRefQualifier());
                    }

                    _writer.WriteIndentedLine("}");
                }

                _writer.WriteIndentedLine("}");
            }
            else
            {
                _writer.WriteIndented("var temp = ");
                WriteFetchCall(method, containingType, kind, parameters);
                _writer.WriteIndentedLine("CopyResults(keys, results.Span, temp);");
            }
        }

        _writer.WriteIndentedLine("}");

        if (kind is DataLoaderKind.Cache)
        {
            return;
        }

        _writer.WriteLine();
        _writer.WriteIndentedLine("private void CopyResults(");
        using (_writer.IncreaseIndent())
        {
            _writer.WriteIndentedLine(
                "global::{0}<{1}> keys,",
                WellKnownTypes.ReadOnlyList,
                key.ToFullyQualifiedWithNullRefQualifier());
            _writer.WriteIndentedLine(
                "global::{0}<{1}<{2}>> results,",
                WellKnownTypes.Span,
                WellKnownTypes.Result,
                kind is DataLoaderKind.Group
                    ? $"{value.ToClassNonNullableFullyQualifiedWithNullRefQualifier()}[]?"
                    : value.ToNullableFullyQualifiedWithNullRefQualifier());
            _writer.WriteIndentedLine(
                "{0} resultMap)",
                ExtractMapType(method.ReturnType).ToFullyQualifiedWithNullRefQualifier());
        }

        _writer.WriteIndentedLine("{");
        using (_writer.IncreaseIndent())
        {
            _writer.WriteIndentedLine("for (var i = 0; i < keys.Count; i++)");
            _writer.WriteIndentedLine("{");
            using (_writer.IncreaseIndent())
            {
                _writer.WriteIndentedLine("var key = keys[i];");
                if (kind is DataLoaderKind.Group)
                {
                    _writer.WriteIndentedLine("if (resultMap.Contains(key))");
                    _writer.WriteIndentedLine("{");

                    using (_writer.IncreaseIndent())
                    {
                        _writer.WriteIndentedLine(
                            "var items = resultMap[key];");
                        _writer.WriteIndentedLine(
                            "results[i] = global::{0}<{1}{2}[]?>.Resolve(global::{3}.ToArray(items));",
                            WellKnownTypes.Result,
                            value.ToFullyQualified(),
                            value.PrintNullRefQualifier(),
                            WellKnownTypes.EnumerableExtensions);
                    }

                    _writer.WriteIndentedLine("}");
                    _writer.WriteIndentedLine("else");
                    _writer.WriteIndentedLine("{");

                    using (_writer.IncreaseIndent())
                    {
                        _writer.WriteIndentedLine(
                            "results[i] = global::{0}<{1}{2}[]?>.Resolve(global::{3}.Empty<{1}{2}>());",
                            WellKnownTypes.Result,
                            value.ToFullyQualified(),
                            value.PrintNullRefQualifier(),
                            WellKnownTypes.Array);
                    }

                    _writer.WriteIndentedLine("}");
                }
                else
                {
                    _writer.WriteIndentedLine("if (resultMap.TryGetValue(key, out var value))");
                    _writer.WriteIndentedLine("{");

                    using (_writer.IncreaseIndent())
                    {
                        _writer.WriteIndentedLine(
                            "results[i] = global::{0}<{1}>.Resolve(value);",
                            WellKnownTypes.Result,
                            value.ToNullableFullyQualifiedWithNullRefQualifier());
                    }

                    _writer.WriteIndentedLine("}");
                    _writer.WriteIndentedLine("else");
                    _writer.WriteIndentedLine("{");

                    using (_writer.IncreaseIndent())
                    {
                        _writer.WriteIndentedLine(
                            "results[i] = global::{0}<{1}>.Resolve(default({2}));",
                            WellKnownTypes.Result,
                            value.ToNullableFullyQualifiedWithNullRefQualifier(),
                            value.IsValueType
                                ? value.ToNullableFullyQualifiedWithNullRefQualifier()
                                : value.ToFullyQualified());
                    }

                    _writer.WriteIndentedLine("}");
                }
            }

            _writer.WriteIndentedLine("}");
        }

        _writer.WriteIndentedLine("}");
    }

    private void WriteFetchCall(
        IMethodSymbol fetchMethod,
        string containingType,
        DataLoaderKind kind,
        ImmutableArray<DataLoaderParameterInfo> parameters)
    {
        _writer.Write("await global::{0}.{1}(", containingType, fetchMethod.Name);

        for (var i = 0; i < parameters.Length; i++)
        {
            if (i > 0)
            {
                _writer.Write(", ");
            }

            var parameter = parameters[i];

            if (i == 0)
            {
                _writer.Write(
                    kind is DataLoaderKind.Cache
                        ? "key"
                        : "keys");
            }
            else
            {
                _writer.Write(parameter.VariableName);
            }
        }

        _writer.WriteLine(").ConfigureAwait(false);");
    }

    public void WriteDataLoaderGroupClass(
        string groupClassName,
        IReadOnlyList<GroupedDataLoaderInfo> dataLoaders,
        bool withInterface,
        bool isInterfacePublic,
        bool isClassPublic)
    {
        if (withInterface)
        {
            _writer.WriteIndentedLine(
                "{0} interface I{1}",
                isInterfacePublic ? "public" : "internal",
                groupClassName);
            _writer.WriteIndentedLine("{");
            _writer.IncreaseIndent();

            foreach (var dataLoader in dataLoaders)
            {
                _writer.WriteIndentedLine("{0} {1} {{ get; }}", dataLoader.InterfaceName, dataLoader.Name);
            }

            _writer.DecreaseIndent();
            _writer.WriteIndentedLine("}");
            _writer.WriteLine();
        }

        if (withInterface)
        {
            _writer.WriteIndentedLine(
                "{0} sealed partial class {1} : I{1}",
                isClassPublic ? "public" : "internal",
                groupClassName);
        }
        else
        {
            _writer.WriteIndentedLine(
                "{0} sealed partial class {1}",
                isClassPublic ? "public" : "internal",
                groupClassName);
        }

        _writer.WriteIndentedLine("{");
        _writer.IncreaseIndent();

        _writer.WriteIndentedLine("private readonly IServiceProvider _services;");

        foreach (var dataLoader in dataLoaders)
        {
            _writer.WriteIndentedLine("private {0}? {1};", dataLoader.InterfaceName, dataLoader.FieldName);
        }

        _writer.WriteLine();
        _writer.WriteIndentedLine("public {0}(IServiceProvider services)", groupClassName);
        _writer.WriteIndentedLine("{");
        _writer.IncreaseIndent();
        _writer.WriteIndentedLine("_services = services");
        _writer.IncreaseIndent();
        _writer.WriteIndentedLine("?? throw new ArgumentNullException(nameof(services));");
        _writer.DecreaseIndent();
        _writer.DecreaseIndent();
        _writer.WriteIndentedLine("}");

        foreach (var dataLoader in dataLoaders)
        {
            _writer.WriteIndentedLine(
                "public {0} {1}",
                dataLoader.InterfaceName,
                dataLoader.Name);

            _writer.WriteIndentedLine("{");
            _writer.IncreaseIndent();

            _writer.WriteIndentedLine("get");

            _writer.WriteIndentedLine("{");
            _writer.IncreaseIndent();

            _writer.WriteIndentedLine(
                "if ({0} is null)",
                dataLoader.FieldName);

            _writer.WriteIndentedLine("{");
            _writer.IncreaseIndent();

            _writer.WriteIndentedLine(
                "{0} = _services.GetRequiredService<{1}>();",
                dataLoader.FieldName,
                dataLoader.InterfaceName);

            _writer.DecreaseIndent();
            _writer.WriteIndentedLine("}");
            _writer.WriteLine();
            _writer.WriteIndentedLine("return {0}!;", dataLoader.FieldName);

            _writer.DecreaseIndent();
            _writer.WriteIndentedLine("}");

            _writer.DecreaseIndent();
            _writer.WriteIndentedLine("}");
        }

        _writer.DecreaseIndent();
        _writer.WriteIndentedLine("}");
    }

    public void WriteLine() => _writer.WriteLine();

    private static ITypeSymbol ExtractMapType(ITypeSymbol returnType)
    {
        if (returnType is INamedTypeSymbol { TypeArguments.Length: 1 } namedType
            && namedType.TypeArguments[0] is INamedTypeSymbol { TypeArguments.Length: 2 } dict)
        {
            return dict;
        }

        throw new InvalidOperationException();
    }

    public override string ToString()
        => _sb.ToString();

    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }

        PooledObjects.Return(_sb);
        _sb = null!;
        _writer = null!;
        _disposed = true;
    }
}
